import torch
import torch.nn as nn
import torch.nn.functional as F

class SimplifiedVGGFeatureExtractor(nn.Module):
    def __init__(self):
        super(SimplifiedVGGFeatureExtractor, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1), 
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 32, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(32, 64, kernel_size=3, padding=1),  
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(64, 128, kernel_size=3, padding=1), 
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        
        self.avgpool = nn.AdaptiveAvgPool2d((4, 4))  
        self.flatten = nn.Flatten()

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = self.flatten(x)
        return x

class MixtureOfExperts(nn.Module):
    def __init__(self, num_experts):
        super(MixtureOfExperts, self).__init__()
        
        self.vgg_extractor = SimplifiedVGGFeatureExtractor()
        vgg_output_size = 128 * 4 * 4  
        self.expert_selector = nn.Linear(vgg_output_size, num_experts) 
        self.experts = nn.ModuleList([nn.Linear(vgg_output_size, vgg_output_size) for _ in range(num_experts)])  # 专家网络

    def forward(self, x):
        features = self.vgg_extractor(x)
        gating_logits = self.expert_selector(features)
        gating_weights = F.softmax(gating_logits, dim=1) 
        expert_outputs = [expert(features) for expert in self.experts]
        expert_outputs = torch.stack(expert_outputs, dim=1)  
        
        final_output = torch.sum(gating_weights.unsqueeze(2) * expert_outputs, dim=1)  
        
        return final_output
